import torch
from torchvision import transforms
from PIL import Image, ImageOps
from torch.utils.data import DataLoader

import numpy as np
import scipy.misc as misc
import os
import glob
import csv
from dataloader.load_DR import Data
from utils.build_dataset import write2csv_pred
from utils.post_process import get_class_acc, get_class_acc_V2, plot_ROC
from compute_evaluate_metrics import plot_performance
from model.fine_tune import BiResNet
seed = 0
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
# torch.backends.cudnn.determinstic = True
# torch.backends.cudnn.benchmark = False

# DATABASE = 'TNT/'
DATABASE = '/media/imed/Data/Experiments/DR_lei/'
# DATABASE = 'cleanedData/'
MODEL_NAME = "DR2"
#
args = {
    'root'      : '/media/imed/Data/Experiments/DR_lei/',
    'pred_path' : "./assets/",
    'model_name': "DR2"
}
def load_net(net):
    # torch.nn.Module.dump_patches = True
    net.load_state_dict(torch.load("/home/imed/Research/TortuosityGrading/checkpoint/DR_grading.pkl", map_location='cpu'))
    # net = torch.load("/home/imed/Research/TortuosityGrading/checkpoint/BANet34-DR.pth")
    print(net)
    return net
def predict():
    net = BiResNet(num_class=5, model_name='resnet34')
    net = load_net(net)

    data = Data(root_dir=args['root'], train=False)
    test_data = DataLoader(data, batch_size=1, shuffle=True, num_workers=2)
    with torch.no_grad():
        net.eval()
        print("Predicting ...")
        for idx, batch in enumerate(test_data):
            img = batch[0]["img"]
            label = batch[1]["img_id"]

            x1, x2, predictions = net(img)
            # x1, x2, predictions = net(image)

            probs = torch.softmax(predictions, dim=1)
            probs = probs.data.cpu().numpy()
            predictions = torch.argmax(predictions, dim=1)
            predictions = predictions.data.cpu().numpy()
            for i in range(len(predictions)):
                file_name = os.path.basename(batch[1]["img_name"][i])
                class_name = predictions[i]
                p = probs[i][predictions[i]]
                p2 = probs[i]
                p2 = ['{:.4f}'.format(i) for i in p2]
                save_folder = os.path.join(args['pred_path'], MODEL_NAME)
                if not os.path.exists(save_folder):
                    os.makedirs(save_folder)
                write2csv_pred(os.path.join(save_folder, "result_pred.csv"), file_name, class_name, p)
                # write2csv_pred(os.path.join(save_folder, "result_pred2222.csv"), file_name, class_name, p2)
    print("Plotting results ...")
    avg_acc, avg_f1 = get_class_acc_V2(os.path.join(save_folder, "result_pred.csv"), args['model_name'])
    print("acc:{0:.4f}\nF1:{1:.4f}".format(avg_acc, avg_f1))
    plot_performance(os.path.join(save_folder, "result_pred.csv"))
if __name__ == '__main__':
    predict()
